using Crux.Core.LangChain.Abstractions.Databases;
using Crux.Core.LangChain.Abstractions.Databases.Metadatas;
using Crux.Core.LangChain.Abstractions.DocumentLoaders;
using Crux.Core.LangChain.Abstractions.Embeddings;

namespace Crux.Core.LangChain.Abstractions.Extensions;

/// <summary>
/// 
/// </summary>
public static class VectorCollectionExtensions
{
    /// <summary>
    /// 
    /// </summary>
    /// <param name="documents"></param>
    /// <param name="separator"></param>
    /// <returns></returns>
    public static string AsString(
        this IEnumerable<Document> documents,
        string separator = "\n\n")
    {
        return string.Join(separator, documents.Select(x => x.PageContent));
    }
    /// <summary>
    /// 
    /// </summary>
    /// <param name="vectorCollection"></param>
    /// <param name="service"></param>
    /// <param name="prompt"></param>
    /// <param name="dimensions"></param>
    /// <param name="amount"></param>
    /// <param name="cancellationToken"></param>
    /// <returns></returns>
    /// <exception cref="ArgumentNullException"></exception>
    public static async Task<IReadOnlyCollection<Document>> GetSimilarDocuments(
        this IVectorCollection vectorCollection,
        ITextEmbeddingGenerationService service,
        string prompt,
        int dimensions = 1536,
        int amount = 4,
        CancellationToken cancellationToken = default)
    {
        if (service is null)
        {
            throw new ArgumentNullException(nameof(service));
        }

        var embedding = await service.GenerateEmbeddingAsync(prompt, dimensions, cancellationToken).ConfigureAwait(false);

        var setting = new VectorSearchSettings
        {
            NumberOfResults = amount
        };

        var vectors = await vectorCollection.SearchAsync(embedding, setting, cancellationToken).ConfigureAwait(false);

        return vectors.Items
            .Select(x => new Document(x.Text, x.Metadata))
            .ToArray();
    }
}